"""
任务调度器模块
基于Celery Beat实现定时任务和语义相似度任务预调度
"""

import logging
from typing import Dict, Any, List, Optional
from datetime import datetime, timedelta
import asyncio
import uuid

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, func, desc
from loguru import logger

from app.core.database import AgentTask, KnowledgeBase, UserInteraction, get_db
from app.core.db_utils import db_utils
from middleware.redis_adapter import get_redis_adapter
from states.celery_manager import get_celery_manager
from utils.constants import TaskPriority, TaskStatus


class TaskScheduler:
    """任务调度器
    
    功能:
    1. 定时任务管理（知识库更新、日志清理等）
    2. 基于语义相似度的任务预调度
    3. 任务优先级管理
    4. 智能体负载均衡
    """
    
    def __init__(self):
        self.redis_client = get_redis_adapter()
        self.celery_manager = get_celery_manager()
        self.logger = logger.bind(module="task_scheduler")
        self._initialize_scheduled_tasks()
    
    def _initialize_scheduled_tasks(self):
        """初始化定时任务配置"""
        # 这些配置会在Celery Beat中生效
        # 实际的任务注册在celery_app.py中
        self.scheduled_tasks = {
            "update_knowledge_index": {
                "schedule": timedelta(minutes=30),
                "task": "app.core.task_scheduler.update_knowledge_index",
                "args": ()
            },
            "cleanup_old_logs": {
                "schedule": timedelta(hours=24),
                "task": "app.core.task_scheduler.cleanup_old_logs",
                "args": (7,)  # 保留7天日志
            },
            "optimize_task_queue": {
                "schedule": timedelta(minutes=15),
                "task": "app.core.task_scheduler.optimize_task_queue",
                "args": ()
            },
            "prewarm_models": {
                "schedule": timedelta(hours=1),
                "task": "app.core.task_scheduler.prewarm_models",
                "args": ()
            }
        }
    
    async def schedule_task(self, task_data: Dict[str, Any]) -> str:
        """
        调度任务到合适的智能体
        
        Args:
            task_data: 任务数据，包含task_type, input_data, priority等
            
        Returns:
            task_id: 调度后的任务ID
        """
        # 生成任务ID
        task_id = str(uuid.uuid4())
        
        # 基于语义相似度的任务预调度
        best_agent = await self._find_best_agent_for_task(task_data)
        
        # 设置任务优先级
        priority = await self._calculate_task_priority(task_data)
        
        # 保存任务到数据库
        task = AgentTask(
            id=task_id,
            task_type=task_data.get("task_type", "unknown"),
            status=TaskStatus.PENDING.value,
            input_data=task_data,
            agent_id=best_agent,
            priority=priority
        )
        
        async for db in get_db():
            try:
                db.add(task)
                await db.commit()
                await db.refresh(task)
                break
            except Exception as e:
                await db.rollback()
                self.logger.error(f"保存任务失败: {str(e)}")
                raise
        
        # 提交到Celery队列
        await self._submit_to_celery(task_id, best_agent, task_data, priority)
        
        return task_id
    
    async def _find_best_agent_for_task(self, task_data: Dict[str, Any]) -> str:
        """
        基于语义相似度查找最适合执行任务的智能体
        
        Args:
            task_data: 任务数据
            
        Returns:
            agent_id: 最适合的智能体ID
        """
        try:
            # 获取任务输入文本
            input_text = self._extract_text_from_task(task_data)
            if not input_text:
                # 默认返回文本问答智能体
                return "text_qa_agent"
            
            # 计算任务输入的嵌入向量
            embedding = await self._get_text_embedding(input_text)
            
            # 查询历史相似任务的执行情况
            similar_agents = await self._get_similar_task_agents(embedding, task_data.get("task_type"))
            
            if similar_agents:
                # 选择成功率最高的智能体
                return max(similar_agents, key=lambda x: x[1])[0]
            
            # 基于任务类型的默认智能体映射
            agent_mapping = {
                "text_qa": "text_qa_agent",
                "multi_modal": "multi_modal_agent",
                "doc_segment": "doc_segment_agent",
                "tool_call": "tool_call_agent"
            }
            
            return agent_mapping.get(task_data.get("task_type"), "text_qa_agent")
            
        except Exception as e:
            self.logger.error(f"查找最佳智能体失败: {str(e)}")
            # 降级到默认智能体
            return "text_qa_agent"
    
    def _extract_text_from_task(self, task_data: Dict[str, Any]) -> Optional[str]:
        """从任务数据中提取文本内容"""
        # 尝试从不同的字段获取文本
        text_fields = ["query", "input", "question", "content", "text"]
        
        for field in text_fields:
            if field in task_data:
                return str(task_data[field])
        
        # 如果input_data是字典，尝试从中提取
        if isinstance(task_data.get("input_data"), dict):
            for field in text_fields:
                if field in task_data["input_data"]:
                    return str(task_data["input_data"][field])
        
        return None
    
    async def _get_text_embedding(self, text: str) -> List[float]:
        """获取文本的嵌入向量"""
        # 使用现有的知识库工具获取嵌入
        try:
            # 这里调用db_utils中的嵌入生成功能
            embedding = await db_utils.generate_embedding(text)
            return embedding
        except Exception as e:
            self.logger.error(f"生成嵌入向量失败: {str(e)}")
            # 返回默认向量
            return [0.0] * 1536  # 假设使用1536维向量
    
    async def _get_similar_task_agents(self, embedding: List[float], task_type: Optional[str] = None) -> List[tuple]:
        """获取相似任务的智能体执行情况"""
        async for db in get_db():
            try:
                # 构建查询
                query = select(
                    AgentTask.agent_id,
                    func.avg(
                        func.case(
                            (AgentTask.status == TaskStatus.COMPLETED.value, 1),
                            else_=0
                        )
                    ).label('success_rate'),
                    func.count(AgentTask.id).label('task_count')
                )
                
                # 添加任务类型过滤
                if task_type:
                    query = query.where(AgentTask.task_type == task_type)
                
                # 只考虑已完成或失败的任务
                query = query.where(
                    AgentTask.status.in_([TaskStatus.COMPLETED.value, TaskStatus.FAILED.value])
                )
                
                # 按智能体分组并排序
                query = query.group_by(AgentTask.agent_id).order_by(desc('success_rate'))
                
                result = await db.execute(query)
                return result.all()
                
            except Exception as e:
                self.logger.error(f"查询相似任务失败: {str(e)}")
                return []
    
    async def _calculate_task_priority(self, task_data: Dict[str, Any]) -> int:
        """计算任务优先级"""
        # 基础优先级
        base_priority = TaskPriority.NORMAL.value
        
        # 根据任务类型调整优先级
        task_type_priority = {
            "urgent": TaskPriority.HIGH.value,
            "critical": TaskPriority.URGENT.value
        }
        
        # 检查是否有明确的优先级设置
        if "priority" in task_data:
            if isinstance(task_data["priority"], str):
                return task_type_priority.get(task_data["priority"], base_priority)
            elif isinstance(task_data["priority"], int):
                return max(0, min(10, task_data["priority"]))  # 确保在0-10范围内
        
        # 根据任务内容的紧急程度动态调整
        urgent_keywords = ["紧急", "urgent", "立即", "immediate", "critical"]
        input_text = self._extract_text_from_task(task_data)
        
        if input_text:
            input_lower = input_text.lower()
            for keyword in urgent_keywords:
                if keyword.lower() in input_lower:
                    return TaskPriority.HIGH.value
        
        return base_priority
    
    async def _submit_to_celery(self, task_id: str, agent_id: str, task_data: Dict[str, Any], priority: int):
        """将任务提交到Celery队列"""
        try:
            # 构建任务参数
            task_params = {
                "task_id": task_id,
                "agent_id": agent_id,
                "task_data": task_data,
                "priority": priority
            }
            
            # 提交任务到Celery
            await self.celery_manager.submit_task(
                task_name="execute_agent_task",
                args=[],
                kwargs=task_params,
                queue=agent_id.replace("_", "."),  # 假设队列名为智能体ID的点分隔形式
                priority=priority
            )
            
            self.logger.info(f"任务已提交: task_id={task_id}, agent_id={agent_id}, priority={priority}")
            
        except Exception as e:
            self.logger.error(f"提交Celery任务失败: {str(e)}")
            raise
    
    async def update_knowledge_index(self):
        """更新知识库索引"""
        self.logger.info("开始更新知识库索引...")
        try:
            # 这里实现知识库索引更新逻辑
            # 例如：合并增量更新、优化向量索引等
            
            # 模拟索引更新过程
            await asyncio.sleep(5)
            
            self.logger.info("知识库索引更新完成")
            
        except Exception as e:
            self.logger.error(f"更新知识库索引失败: {str(e)}")
            raise
    
    async def cleanup_old_logs(self, days: int = 7):
        """清理旧日志"""
        self.logger.info(f"开始清理{days}天前的日志...")
        try:
            # 计算截止日期
            cutoff_date = datetime.utcnow() - timedelta(days=days)
            
            # 清理数据库中的旧交互记录
            async for db in get_db():
                try:
                    # 删除旧的用户交互记录
                    await db.execute(
                        UserInteraction.__table__.delete().where(
                            UserInteraction.created_at < cutoff_date
                        )
                    )
                    await db.commit()
                    break
                except Exception as e:
                    await db.rollback()
                    self.logger.error(f"清理数据库日志失败: {str(e)}")
                    raise
            
            self.logger.info(f"日志清理完成，已删除{days}天前的记录")
            
        except Exception as e:
            self.logger.error(f"清理日志失败: {str(e)}")
            raise
    
    async def optimize_task_queue(self):
        """优化任务队列"""
        self.logger.info("开始优化任务队列...")
        try:
            # 获取队列统计信息
            queue_stats = await self.celery_manager.get_queue_stats()
            
            # 根据队列长度动态调整任务优先级
            for queue_name, stats in queue_stats.items():
                if stats.get("task_count", 0) > 100:  # 如果队列积压严重
                    # 实现队列优化逻辑
                    self.logger.info(f"队列{queue_name}积压严重，优化中...")
            
            self.logger.info("任务队列优化完成")
            
        except Exception as e:
            self.logger.error(f"优化任务队列失败: {str(e)}")
    
    async def prewarm_models(self):
        """预加载模型"""
        self.logger.info("开始预加载模型...")
        try:
            # 提交预热任务到各个智能体
            for agent_id in ["text_qa_agent", "multi_modal_agent", "doc_segment_agent"]:
                prewarm_task_id = f"prewarm_{str(uuid.uuid4())[:8]}"
                await self.celery_manager.submit_task(
                    task_name="prewarm_model",
                    args=[],
                    kwargs={"agent_id": agent_id},
                    queue=agent_id.replace("_", ".")
                )
                self.logger.info(f"已提交模型预热任务: {prewarm_task_id} for {agent_id}")
            
            self.logger.info("模型预加载完成")
            
        except Exception as e:
            self.logger.error(f"预加载模型失败: {str(e)}")
    
    async def predict_peak_load(self, hours_ahead: int = 1):
        """
        预测未来负载高峰，用于任务预调度
        
        Args:
            hours_ahead: 预测未来几小时
            
        Returns:
            负载预测结果
        """
        try:
            # 这里可以实现更复杂的负载预测算法
            # 例如：基于历史数据的时间序列预测
            
            # 获取历史负载数据
            async for db in get_db():
                result = await db.execute(
                    select(
                        func.date_trunc('hour', AgentTask.created_at).label('hour'),
                        func.count(AgentTask.id).label('task_count')
                    ).where(
                        AgentTask.created_at >= datetime.utcnow() - timedelta(days=7)
                    ).group_by(
                        func.date_trunc('hour', AgentTask.created_at)
                    ).order_by('hour')
                )
                history_data = result.all()
                break
            
            # 简单的时间模式预测
            # 这里可以替换为更复杂的预测模型
            hour_of_day = (datetime.utcnow().hour + hours_ahead) % 24
            
            # 统计历史上相同小时的平均负载
            same_hour_tasks = [h.task_count for h in history_data if h.hour.hour == hour_of_day]
            avg_load = sum(same_hour_tasks) / len(same_hour_tasks) if same_hour_tasks else 0
            
            prediction = {
                "hour": hour_of_day,
                "predicted_load": avg_load,
                "is_peak": avg_load > 50,  # 假设50个任务/小时为高峰
                "recommended_action": "prewarm" if avg_load > 50 else "normal"
            }
            
            return prediction
            
        except Exception as e:
            self.logger.error(f"预测负载失败: {str(e)}")
            # 返回默认预测
            return {
                "hour": (datetime.utcnow().hour + hours_ahead) % 24,
                "predicted_load": 0,
                "is_peak": False,
                "recommended_action": "normal"
            }


# 创建全局任务调度器实例
task_scheduler = TaskScheduler()


# Celery任务定义（这些函数将被注册到Celery中）
def update_knowledge_index_task():
    """Celery任务：更新知识库索引"""
    asyncio.run(task_scheduler.update_knowledge_index())


def cleanup_old_logs_task(days: int = 7):
    """Celery任务：清理旧日志"""
    asyncio.run(task_scheduler.cleanup_old_logs(days))


def optimize_task_queue_task():
    """Celery任务：优化任务队列"""
    asyncio.run(task_scheduler.optimize_task_queue())


def prewarm_models_task():
    """Celery任务：预加载模型"""
    asyncio.run(task_scheduler.prewarm_models())


def execute_agent_task(task_id: str, agent_id: str, task_data: Dict[str, Any], priority: int):
    """Celery任务：执行智能体任务"""
    # 这里将通过异步编排器执行实际任务
    # 由于Celery是同步的，需要使用asyncio.run来执行异步代码
    from apps.orchestrator.async_orchestrator_agent import get_async_orchestrator
    
    async def execute_task():
        orchestrator = await get_async_orchestrator()
        return await orchestrator.execute_task_async(
            task_id=task_id,
            agent_id=agent_id,
            task_data=task_data,
            priority=priority
        )
    
    return asyncio.run(execute_task())


def prewarm_model_task(agent_id: str):
    """Celery任务：预加载特定智能体的模型"""
    async def prewarm():
        # 实现模型预加载逻辑
        logger.info(f"预加载智能体{agent_id}的模型")
        # 这里可以调用智能体的预热方法
        await asyncio.sleep(2)
        return {"status": "success", "agent_id": agent_id}
    
    return asyncio.run(prewarm())
